from typing import Annotated, Sequence, Literal
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables.graph import MermaidDrawMethod
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import ToolNode, tools_condition
from typing_extensions import TypedDict
from langchain_community.vectorstores import Chroma
from langchain_core.tools import create_retriever_tool

from langchain_core.messages import BaseMessage, HumanMessage
from langgraph.graph import add_messages, StateGraph,END,START

from pydantic import BaseModel, Field

from qw_model import qw_model
from Embed_model import BGE_Embed

# -------------------------------------定义工具
# 加载知识库
embedding = BGE_Embed()
vectorstore = Chroma(
    embedding_function=embedding,
    persist_directory='./chroma_db'
)

# 创建检索器
retriever = vectorstore.as_retriever(search_kwargs={'k': 3})

# 创建retriever工具
retriever_tool = create_retriever_tool(
    retriever,
    "yingong",
    "业务咨询和价格查询"
)

# 工具列表
tools = [retriever_tool]

# # 检索出的文档
# result = retriever.invoke('抖音云剪多少钱')
# print(result)


# ————————————————————————————————————定义节点
# 状态图
class AgentState(TypedDict):
    # add_messages 函数定义了应如何处理更新。
    # 默认是替换。add_messages 表示 “追加”。
    messages: Annotated[Sequence[BaseMessage], add_messages]

# 定义检索的条件边函数
def grade_documents(state) ->Literal["generate","rewrite"]:
    """
    确定检索到的文档是否与问题相关。
    Args：
    state（消息）：当前状态
    Returns：
    str：决定文件是否相关
    """
    class grade(BaseModel):
        """相关性检查的二进制分数"""
        binary_score: str = Field(description="相关性分数 'yes' or 'no'")

    #LLM
    model = qw_model()
    # 具有工具和验证的大语言模型
    llm_with_tool = model.with_structured_output(grade)

    #提示词
    prompt = PromptTemplate(
        template="""你是评估检索到的文档与用户问题的相关性的评分者。\n 
         这是检索到的文档: \n\n {context} \n\n
         这是用户问题: {question} \n
         如果文档包含与用户问题相关的关键字或语义，请将其评为相关。 \n
         给出一个二进制分数 "yes" 或 "no" 分数来指示文档是否与问题相关。""",
        input_variables=["context", "question"],
    )

    # Chain
    chain = prompt | llm_with_tool

    messages = state["messages"]
    last_message = messages[-1]

    question = messages[0].content
    docs = last_message.content

    scored_result = chain.invoke({"question": question, "context": docs})

    # 得分
    score = scored_result.binary_score

    if score == "yes":
        # print("---文档相关---")
        return "generate"

    else:
        print("---文档不相关---")
        print(score)
        return "rewrite"

# 定义调用模型的函数节点
def agent(state):
    """
   调用代理模型以根据当前状态生成响应。给定问题，它将决定使用检索器工具进行检索，还是简单地结束。
   Args：
   state（消息）：当前状态
   Returns：
   file：代理响应附加到消息的更新状态
    """
    messages = state["messages"]
    model = qw_model()
    model = model.bind_tools(tools)
    response = model.invoke(messages)
    # 我们返回一个列表，因为这将被添加到现有列表中
    return {"messages": [response]}

# 定义重写函数节点
def rewrite(state):
    """
    转换查询以产生更好的问题。
    Args：
    state（消息）：当前状态
    Returns：
    结论：带有重新措辞问题的更新状态
    """
    messages = state["messages"]
    question = messages[0].content

    msg = [
        HumanMessage(
            content=f""" \n 
    查看输入并根据上下文信息，尝试推理潜在的语义意图 / 含义。 \n 
    这是最初的问题:
    \n ------- \n
    {question} 
    \n ------- \n
    提出一个改进的问题: """,
        )
    ]

    model =qw_model()
    response = model.invoke(msg)
    return {"messages": [response]}

def generate(state):
    """
   生成答案
   Args：
   state（消息）：当前状态
   Returns：
   结论：带有重新措辞问题的更新状态
    """
    messages = state["messages"]
    question = messages[0].content
    last_message = messages[-1]

    docs = last_message.content

    # Prompt
    prompt = PromptTemplate(
        template="""你是一个聪明又有亲和力的销售客服，名叫银弓，擅长用轻松、亲切的语气解答用户关于公司业务、价格、产品咨询的问题。你的目标是提供专业、准确的回答，尤其是价格问题，严格按照相应文档回答，风格需要口语化、自然流畅、贴近用户生活，避免书面化和过于正式的表达以及尊称。\n 
             这是检索到的文档: \n\n {context} \n\n
             这是用户问题: {question} \n
             "1. 当用户问题不清楚时，主动询问用户具体需求。"
             "2. 回答价格问题，严格按照检索的文档给出，不可自行编造。"
             "3. 当没有用户所说的地区，就按它的省份价格算，当省份也没有那就只能做不限地区的"
             "4. 保持沟通语气亲切、自然，避免过于正式的措辞，多用口语化的表达和轻松的语气，比如“老板”、“宝子”、“咱们”，禁止使用'您'。"
             "5. 输出答案应简洁明了，字数在15字左右。" \n
              需满足以上回答要求。""",
        input_variables=["context", "question"],
    )

    # LLM
    llm = qw_model()

    # Chain
    rag_chain = prompt | llm | StrOutputParser()

    response = rag_chain.invoke({"context": docs, "question": question})
    return {"messages": [response]}

# print("*" * 20 + "Prompt[rlm/rag-prompt]" + "*" * 20)
# prompt = hub.pull("rlm/rag-prompt").pretty_print()

#---------------------------------------------初始化图形
# 定义转状态图
w = StateGraph(AgentState)

# 节点
w.add_node("agent",agent)
retrieve = ToolNode([retriever_tool])
w.add_node("retrieve",retrieve)
w.add_node("rewrite",rewrite)
w.add_node("generate",generate)

w.add_edge(START, "agent")
# agent条件边
w.add_conditional_edges(
    "agent",
    tools_condition,
    {
        "tools": "retrieve",
        END: END,
    },
)
#检索节点条件边
w.add_conditional_edges(
    "retrieve",
    grade_documents,
    {
        "generate": "generate",
        "rewrite": "rewrite",
    },
)

w.add_edge("rewrite","agent")
w.add_edge("generate", END)

# --------------------------------------------编译图
memory = MemorySaver()
app = w.compile(checkpointer=memory)

config = {"configurable": {'thread_id': "a1"}}

# # 可视化
# graph_img = app.get_graph().draw_mermaid_png(
#     draw_method=MermaidDrawMethod.API,
# )
# with open("graph_1.png", "wb") as f:
#     f.write(graph_img)


res = app.invoke({
    "messages": [
        ("user", "我叫张说"),
    ]
}, config)
print(res["messages"][-1].content)

res2 = app.invoke({
    "messages": [
        ("user", "我叫什么"),
    ]
}, config)
print(res2["messages"][-1].content)